/* benchmark guru interfaces */
#include <Array.h>
#include <cstdlib>
#include <complex>
#include <iostream>
#include <iomanip>
using namespace std;

template<typename T>
struct ArrayProxy
{
  T* data_ptr;
  size_t Length[3];
  ArrayProxy(T* in, size_t n0, size_t n1, size_t  n2): data_ptr(in)
  {
    Length[0]=n0; Length[1]=n1; Length[2]=n2;
  }
  inline size_t size(int i) const { return Length[i];}
  inline T operator()(int i, int j, int k) const { return data_ptr[k+Length[2]*(j+i*Length[1])];}
  inline T& operator()(int i, int j, int k) { return data_ptr[k+Length[2]*(j+i*Length[1])];}
};

template<typename T>
inline void transpose_yz(const Array<T,3>& in, Array<T,3>& out, ArrayProxy<T>& temp)
{
  for(int k=0; k<temp.size(0);++k)
    for(int i=0; i<temp.size(1); ++i)
      for(int j=0; j<temp.size(2);++j)
        temp(k,i,j)=in(i,j,k);

  //alltoall
  
  for(int i=0; i<out.size(0);++i)
    for(int k=0; k<out.size(1); ++k)
      for(int j=0; j<out.size(2);++j)
        out(i,k,j)=temp(k,i,j);
}

template<typename T>
inline void transpose_xy(const Array<T,3>& in, Array<T,3>& out, ArrayProxy<T>& temp)
{
  for(int k=0; k<temp.size(0);++k)
    for(int i=0; i<temp.size(1); ++i)
      for(int j=0; j<temp.size(2);++j)
        temp(k,i,j)=in(i,j,k);

  //alltoall
  
  for(int k=0; k<out.size(0);++k)
    for(int j=0; j<out.size(1); ++j)
      for(int i=0; i<out.size(2);++i)
        out(k,j,i)=temp(k,i,j);
}

template<typename T>
inline void transpose_yz(const Array<T,3>& in, Array<T,3>& out)
{
  
  for(int i=0; i<out.size(0);++i)
    for(int j=0; j<out.size(1); ++j)
      for(int k=0; k<out.size(2); ++k)
        out(i,j,k)=in(i,k,j);
}

template<typename T>
inline void transpose_xy(const Array<T,3>& in, Array<T,3>& out)
{
  for(int i=0; i<out.size(0);++i)
    for(int j=0; j<out.size(1); ++j)
      for(int k=0; k<out.size(2); ++k)
        out(i,j,k)=in(k,j,i);
}

int main(int argc, char** argv)
{
  //mpi::environment env(argc,argv);
  //mpi::communicator mycomm;
  //OhmmsInfo ohmms("transpose",mycomm.rank(),0,1);

  ///number of iterations
  int niters=10;
  const int APP_DIM=3;
  const int default_fft_dim=4;
  const int ndim = 2;
  // default dimension 128 128 128  with 2D distribution
  int nx=default_fft_dim, ny=default_fft_dim, nz=default_fft_dim;

  bool debug=false;
  int npx=1,npy=1,np=1;
  int ic=0;
  while(ic<argc)
  {
    std::string a(argv[ic]);
    if(a == "-np")
    {
      np=atoi(argv[++ic]);
    }
    else if(a == "-npx")
    {
      npx=atoi(argv[++ic]);
    }
    else if(a == "-s") 
      nx=ny=nz=atoi(argv[++ic]);
    else if(a == "-i")
      niters=atoi(argv[++ic]);
    else if(a == "-debug")
      debug=true;
    ++ic;
  }

  npy=np/npx;

  if(npx*npy != np) 
  {
    cerr << "npx*npy != np " << endl;
    return 1;
  }

  typedef double real_type;
  typedef complex<double> value_type;
  //typedef double value_type;

  Array<value_type,3> Az(nx/npx,ny/npy,nz);
  Array<value_type,3> Ay(nx/npx,nz/npy,ny);
  Array<value_type,3> Ax(ny/npx,nz/npy,nx);

  vector<value_type> temp(Az.size());//this is a dummy array
  ArrayProxy<value_type> tzy(&temp[0],Az.size(2),Az.size(0),Az.size(1));
  ArrayProxy<value_type> tyx(&temp[0],Ay.size(2),Ay.size(0),Ay.size(1));

  {
    int ijk=0;
    for(int i=0; i<Az.size(0);++i)
      for(int j=0; j<Az.size(1);++j)
        for(int k=0; k<Az.size(2);++k)
          Az(ijk)=ijk++;
  }

  transpose_yz(Az,Ay,tzy);
  transpose_xy(Ay,Ax,tyx);

  if(debug)
  {
    for(int i=0; i<Az.size(0);++i)
      for(int j=0; j<Az.size(1);++j)
      {
        cout << setw(4) << i << setw(4) <<j << " : " ;
        for(int k=0; k<Az.size(2);++k) cout << setw(4) << Az(i,j,k) << "(" << Ay(i,k,j) << ") ";
        cout  << endl;
      }

    cout << "----------------------------------------------------- " << endl;
    for(int i=0; i<Az.size(0);++i)
      for(int j=0; j<Az.size(1);++j)
      {
        cout << setw(4) << i << setw(4) <<j << " : " ;
        for(int k=0; k<Az.size(2);++k) cout << setw(4) << Az(i,j,k) << "(" << Ax(j,k,i) << ") ";
        cout  << endl;
      }
    cout << "----------------------------------------------------- " << endl;
    for(int i=0; i<Ay.size(0);++i)
      for(int j=0; j<Ay.size(1);++j)
      {
        cout << setw(4) << i << setw(4) <<j << " : " ;
        for(int k=0; k<Ay.size(2);++k) cout << setw(4) << Ay(i,j,k) << "(" << Ax(k,j,i) << ") ";
        cout  << endl;
      }
  }
  return 0;
}

